#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.time import Time
from rclpy.action import ActionClient
from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import ReentrantCallbackGroup, MutuallyExclusiveCallbackGroup

from geometry_msgs.msg import Twist, Point
from std_msgs.msg import String, Bool
from nav_msgs.msg import Odometry
from ball_collector_interfaces.msg import BallInfo, RobotState
from ball_collector_interfaces.srv import SetGripper, ChangeColor
import threading
import time
import math
import random
from enum import Enum


class TestStateMonitor(Node):
    """
    监控策略节点状态的测试节点
    """
    def __init__(self):
        super().__init__('test_state_monitor')
        
        # 创建互斥回调组
        self.timer_cb_group = MutuallyExclusiveCallbackGroup()
        self.topic_cb_group = MutuallyExclusiveCallbackGroup()
        
        # 状态追踪
        self.current_state = "IDLE"
        self.has_ball = False
        self.cmd_vel_history = []
        self.state_history = []
        self.gripper_calls = []
        self.target_color = "red"
        
        # 标记测试进度
        self.test_step = 0
        self.step_start_time = time.time()
        self.test_completed = False
        
        # 订阅策略节点的输出
        self.cmd_vel_sub = self.create_subscription(
            Twist,
            'cmd_vel',
            self.cmd_vel_callback,
            10,
            callback_group=self.topic_cb_group
        )
        
        self.robot_state_sub = self.create_subscription(
            RobotState,
            'robot_state',
            self.robot_state_callback,
            10,
            callback_group=self.topic_cb_group
        )
        
        # 创建发布者
        self.ball_pub = self.create_publisher(
            BallInfo,
            'detected_ball',
            10
        )
        
        self.odom_pub = self.create_publisher(
            Odometry,
            'odom',
            10
        )
        
        # 创建模拟服务
        self.gripper_service = self.create_service(
            SetGripper,
            '/hardware_interface/set_gripper',
            self.gripper_callback
        )
        
        self.color_service = self.create_service(
            ChangeColor,
            'change_target_color',
            self.color_callback
        )
        
        # 创建定时器
        self.test_timer = self.create_timer(
            0.5,  # 500ms
            self.test_step_callback,
            callback_group=self.timer_cb_group
        )
        
        self.status_timer = self.create_timer(
            5.0,  # 5秒
            self.print_status,
            callback_group=self.timer_cb_group
        )
        
        self.get_logger().info("测试监控节点已初始化")
    
    def cmd_vel_callback(self, msg):
        """记录速度命令"""
        current_time = time.time()
        cmd = {
            'time': current_time,
            'linear_x': msg.linear.x,
            'angular_z': msg.angular.z
        }
        self.cmd_vel_history.append(cmd)
        
        # 仅记录有意义的变化
        if len(self.cmd_vel_history) > 1:
            prev = self.cmd_vel_history[-2]
            if (abs(prev['linear_x'] - cmd['linear_x']) > 0.01 or 
                abs(prev['angular_z'] - cmd['angular_z']) > 0.01):
                self.get_logger().info(f"速度命令: 线速度={msg.linear.x:.2f}, 角速度={msg.angular.z:.2f}")
    
    def robot_state_callback(self, msg):
        """记录机器人状态"""
        if msg.current_state != self.current_state or msg.has_ball != self.has_ball:
            self.get_logger().info(f"状态变化: {self.current_state} -> {msg.current_state}, 持球: {msg.has_ball}")
            
            # 记录状态历史
            self.state_history.append({
                'time': time.time(),
                'state': msg.current_state,
                'has_ball': msg.has_ball
            })
            
            self.current_state = msg.current_state
            self.has_ball = msg.has_ball
            self.target_color = msg.target_color
    
    def gripper_callback(self, request, response):
        """模拟夹爪服务回调"""
        self.get_logger().info(f"夹爪调用: grip={request.grip}, force={request.force}")
        
        # 记录夹爪调用
        self.gripper_calls.append({
            'time': time.time(),
            'grip': request.grip,
            'force': request.force
        })
        
        # 始终返回成功
        response.success = True
        response.message = f"夹爪{'关闭' if request.grip else '打开'}成功"
        return response
    
    def color_callback(self, request, response):
        """模拟颜色更改服务回调"""
        self.get_logger().info(f"颜色更改请求: {request.new_color}")
        
        # 记录目标颜色变化
        self.target_color = request.new_color
        
        # 返回成功
        response.success = True
        response.current_color = request.new_color
        return response
    
    def publish_mock_ball(self, is_visible=True, x=320, y=240, diameter=100):
        """发布模拟球体检测信息"""
        if not is_visible:
            # 不发布，相当于没有检测到球
            return
        
        msg = BallInfo()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "camera_link"
        msg.id = 1
        msg.color = self.target_color
        
        msg.center = Point()
        msg.center.x = float(x)
        msg.center.y = float(y)
        msg.center.z = 0.0
        
        msg.diameter = float(diameter)
        msg.confidence = 0.95
        msg.is_target = True
        
        self.ball_pub.publish(msg)
    
    def publish_mock_odometry(self, x=0.0, y=0.0, theta=0.0):
        """发布模拟里程计数据"""
        msg = Odometry()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = "odom"
        msg.child_frame_id = "base_link"
        
        # 位置
        msg.pose.pose.position.x = x
        msg.pose.pose.position.y = y
        msg.pose.pose.position.z = 0.0
        
        # 将偏航角转换为四元数（简化版本）
        msg.pose.pose.orientation.w = math.cos(theta / 2.0)
        msg.pose.pose.orientation.z = math.sin(theta / 2.0)
        
        self.odom_pub.publish(msg)
    
    def test_step_callback(self):
        """测试步骤回调，模拟不同的测试场景"""
        if self.test_completed:
            return
        
        current_time = time.time()
        elapsed = current_time - self.step_start_time
        
        # 测试步骤
        if self.test_step == 0:
            # 初始化阶段 - 等待策略节点启动
            self.get_logger().info("测试步骤 0: 等待策略节点启动")
            self.publish_mock_odometry(0.0, 0.0, 0.0)
            
            if elapsed > 3.0:
                self.test_step = 1
                self.step_start_time = current_time
        
        elif self.test_step == 1:
            # 测试搜索状态 - 不发送球信息，观察搜索行为
            self.get_logger().info("测试步骤 1: 测试搜索状态")
            self.publish_mock_odometry(0.0, 0.0, elapsed * 0.1)  # 模拟机器人旋转
            
            # 不发布球信息，让策略节点保持搜索
            
            if elapsed > 5.0:
                self.test_step = 2
                self.step_start_time = current_time
        
        elif self.test_step == 2:
            # 测试接近状态 - 发送远处的球信息
            self.get_logger().info("测试步骤 2: 测试接近状态")
            self.publish_mock_odometry(0.0, 0.0, 0.0)
            
            # 发布球信息 - 逐渐增大直径模拟靠近
            ball_diameter = 50 + elapsed * 10  # 直径随时间增长
            self.publish_mock_ball(True, 320, 240, ball_diameter)
            
            if elapsed > 5.0:
                self.test_step = 3
                self.step_start_time = current_time
        
        elif self.test_step == 3:
            # 测试对准状态 - 发送偏离中心的球
            self.get_logger().info("测试步骤 3: 测试对准状态")
            self.publish_mock_odometry(0.5, 0.0, 0.0)  # 模拟向前移动了一些
            
            # 球偏离中心，逐渐移动到中心
            x_pos = 420 - elapsed * 20  # 从右侧逐渐移动到中心
            self.publish_mock_ball(True, x_pos, 240, 150)  # 较大直径表示已经很近
            
            if elapsed > 5.0:
                self.test_step = 4
                self.step_start_time = current_time
        
        elif self.test_step == 4:
            # 测试抓取状态 - 发送中心对准的球
            self.get_logger().info("测试步骤 4: 测试抓取状态")
            self.publish_mock_odometry(1.0, 0.0, 0.0)
            
            # 发布中心球
            self.publish_mock_ball(True, 320, 240, 200)  # 大直径，中心对准
            
            if elapsed > 3.0:
                self.test_step = 5
                self.step_start_time = current_time
        
        elif self.test_step == 5:
            # 测试运输状态 - 不再发送球信息，模拟已抓取
            self.get_logger().info("测试步骤 5: 测试运输状态")
            
            # 模拟移动到目标区域
            x = min(2.0, 1.0 + elapsed * 0.2)
            y = min(1.0, elapsed * 0.2)
            self.publish_mock_odometry(x, y, 0.0)
            
            # 不发布球信息，因为假设已经抓取
            
            if elapsed > 5.0:
                self.test_step = 6
                self.step_start_time = current_time
        
        elif self.test_step == 6:
            # 测试释放状态
            self.get_logger().info("测试步骤 6: 测试释放状态")
            self.publish_mock_odometry(2.0, 1.0, 0.0)  # 在目标位置
            
            if elapsed > 3.0:
                self.test_step = 7
                self.step_start_time = current_time
        
        elif self.test_step == 7:
            # 测试返回状态
            self.get_logger().info("测试步骤 7: 测试返回状态")
            
            # 模拟逐渐返回
            x = max(0.5, 2.0 - elapsed * 0.3)
            y = max(0.0, 1.0 - elapsed * 0.2)
            self.publish_mock_odometry(x, y, elapsed * 0.1)  # 添加一些旋转
            
            if elapsed > 5.0:
                self.test_step = 8
                self.step_start_time = current_time
        
        elif self.test_step == 8:
            # 测试多次循环 - 模拟新的球
            self.get_logger().info("测试步骤 8: 测试新的搜索周期")
            self.publish_mock_odometry(0.5, 0.0, elapsed * 0.1)
            
            # 间歇性发布新球
            if int(elapsed) % 2 == 0:  # 每隔一秒发布
                self.publish_mock_ball(True, 220, 240, 80)  # 较小的球在左侧
            
            if elapsed > 10.0:
                self.test_step = 9
                self.step_start_time = current_time
        
        elif self.test_step == 9:
            # 测试丢失球
            self.get_logger().info("测试步骤 9: 测试球丢失处理")
            self.publish_mock_odometry(1.0, 0.5, 0.0)
            
            # 间歇性发布球，模拟跟踪丢失
            if int(elapsed * 2) % 4 < 1:  # 25%的时间可见
                self.publish_mock_ball(True, 320, 240, 100)
            
            if elapsed > 10.0:
                self.test_step = 10
                self.step_start_time = current_time
        
        elif self.test_step == 10:
            # 测试颜色切换
            self.get_logger().info("测试步骤 10: 测试颜色切换")
            self.publish_mock_odometry(0.0, 0.0, 0.0)
            
            # 测试不同颜色
            colors = ["red", "green", "blue", "yellow"]
            color_index = int(elapsed) % len(colors)
            
            # 只在颜色变化时发布
            if elapsed < 0.5 or int(elapsed) != int(elapsed - 0.5):
                self.publish_mock_ball(True, 320, 240, 100, colors[color_index])
            
            if elapsed > 10.0:
                self.test_step = 11
                self.step_start_time = current_time
        
        elif self.test_step == 11:
            # 测试完成
            self.get_logger().info("测试完成")
            self.test_completed = True
            self.analyze_results()
    
    def print_status(self):
        """打印当前测试状态"""
        if not self.test_completed:
            self.get_logger().info(f"当前测试步骤: {self.test_step}, 当前状态: {self.current_state}")
            
            # 打印最后的命令
            if self.cmd_vel_history:
                last_cmd = self.cmd_vel_history[-1]
                self.get_logger().info(f"最后的速度命令: 线速度={last_cmd['linear_x']:.2f}, 角速度={last_cmd['angular_z']:.2f}")
    
    def analyze_results(self):
        """分析测试结果"""
        self.get_logger().info("测试结果分析:")
        
        # 分析状态转换
        if len(self.state_history) > 1:
            self.get_logger().info("状态转换序列:")
            for i, state in enumerate(self.state_history):
                transition_time = state['time']
                if i > 0:
                    transition_time = state['time'] - self.state_history[i-1]['time']
                self.get_logger().info(f"  {i}: {state['state']} (持球: {state['has_ball']}), 花费时间: {transition_time:.2f}s")
        
        # 分析夹爪调用
        if self.gripper_calls:
            self.get_logger().info("夹爪操作序列:")
            for i, call in enumerate(self.gripper_calls):
                self.get_logger().info(f"  {i}: {'关闭' if call['grip'] else '打开'}, 力度: {call['force']}")
        
        # 检查是否完成了完整的收集周期
        states_seen = set(s['state'] for s in self.state_history)
        expected_states = {"SEARCHING", "APPROACHING", "ALIGNING", "GRASPING", 
                          "TRANSPORTING", "RELEASING", "RETURNING"}
        
        missing_states = expected_states - states_seen
        if missing_states:
            self.get_logger().warn(f"测试未覆盖以下状态: {missing_states}")
        else:
            self.get_logger().info("测试覆盖了所有期望的状态转换")
        
        # 分析是否检测到了球的抓取和释放
        ball_pickup = False
        ball_release = False
        for i in range(1, len(self.state_history)):
            prev = self.state_history[i-1]
            curr = self.state_history[i]
            if not prev['has_ball'] and curr['has_ball']:
                ball_pickup = True
            if prev['has_ball'] and not curr['has_ball']:
                ball_release = True
        
        if ball_pickup:
            self.get_logger().info("✓ 成功检测到球的抓取")
        else:
            self.get_logger().warn("✗ 未检测到球的抓取")
            
        if ball_release:
            self.get_logger().info("✓ 成功检测到球的释放")
        else:
            self.get_logger().warn("✗ 未检测到球的释放")


def main(args=None):
    rclpy.init(args=args)
    
    test_node = TestStateMonitor()
    
    # 创建多线程执行器
    executor = MultiThreadedExecutor()
    executor.add_node(test_node)
    
    # 启动执行器
    executor_thread = threading.Thread(target=executor.spin, daemon=True)
    executor_thread.start()
    
    try:
        # 等待测试完成或超时
        timeout = 120  # 2分钟超时
        start_time = time.time()
        
        while not test_node.test_completed and time.time() - start_time < timeout:
            time.sleep(1.0)
        
        if not test_node.test_completed:
            test_node.get_logger().warn(f"测试超时 ({timeout}秒)")
            test_node.analyze_results()
    
    except KeyboardInterrupt:
        test_node.get_logger().info("测试被用户中断")
    
    finally:
        executor.shutdown()
        test_node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()